# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import abc
import functools

from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.base.book.types import create_market_ticker_event
from coin.exchange.base.ws.handler import WebsocketHandler, _EventDistributor
from coin.flow.topic_record import TopicRecord
from coin.flow.topic_map_util import get_worker_id
from coin.exchange.okex_futures.ws.handler import (OkexFuturesHandler,
                                                   get_product_from_channel_name_cached)


class ArbEventDistributor(_EventDistributor):
  def __init__(self):
    _EventDistributor.__init__(self)

  def publish(self, event_type, key, event):
    _EventDistributor.publish(self, event_type, key, event)


class SwitchingArbEventDistributor(_EventDistributor):
  def __init__(self):
    _EventDistributor.__init__(self)
    self._publish_this = False
    self.publish_try = 0

  def publish(self, event_type, key, event):
    self.publish_try += 1
    if self._publish_this:
      _EventDistributor.publish(self, event_type, key, event)

  def set_publish(self, publish_this):
    self._publish_this = publish_this


class FeedArbitrationHandler(WebsocketHandler):
  EventDistributorType = None
  OldFeedTimeOut = 10 * 1e9

  def __init__(self, *, feed_checker_config=None):
    WebsocketHandler.__init__(self,
                              event_dist=self.EventDistributorType(),
                              feed_checker_config=feed_checker_config)
    self._old_feed_timeout = self.OldFeedTimeOut
    self._hash_last_time = {}
    self._last_feed_ts = {}

  def update_and_check_new(self, feed_timestamp, hash_or_unique_id):
    if hash_or_unique_id not in self._hash_last_time:
      self._hash_last_time[hash_or_unique_id] = feed_timestamp
      return True
    else:
      ret = feed_timestamp - self._hash_last_time[hash_or_unique_id] > self._old_feed_timeout
      self._hash_last_time[hash_or_unique_id] = feed_timestamp
      return ret

  @abc.abstractmethod
  def get_wrapper_callback(self, callback, hash_function):
    raise NotImplementedError()

  def get_worker_last_feed_ts(self, worker_id):
    return self._last_feed_ts.get(worker_id, None)


class MessageLevelFeedArbHandler(FeedArbitrationHandler):
  EventDistributorType = ArbEventDistributor

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._worker_ids = []
    self._handler = None

  def observe_book_builder(self, book_builder_observer):
    assert self._handler is not None
    self._handler.observe_book_builder(book_builder_observer)

  def add_check_product(self, product):
    assert self._handler is not None
    self._handler.add_check_product(product)

  def set_live(self, is_live):
    assert self._handler
    self._handler.set_live(is_live)

  def add_handler(self, worker_id, handler_cls, **kwargs):
    self._worker_ids.append(str(worker_id))
    if self._handler is None:
      self._handler = handler_cls(feed_checker_config=self._feed_checker_config,
                                  event_dist=self._event_dist,
                                  **kwargs)

  def post_subscribe(self, *args, **kawrgs):
    self._handler.post_subscribe(*args, **kawrgs)

  def subscribe_flow(self, flow_sub, topic, callback_attr, hash_function):
    callback = getattr(self._handler, callback_attr)
    for worker_id in self._worker_ids:
      wrapped_callback = functools.partial(self.on_flow_record, callback, hash_function, worker_id)
      flow_sub.subscribe('%s*%s' % (topic, worker_id), wrapped_callback)

  def on_flow_record(self, callback, hash_function, worker_id, record, queue_data, topic_data):
    uid = hash_function(record)
    if self.update_and_check_new(record.timestamp, uid):
      callback(record, queue_data, topic_data)
    self._last_feed_ts[int(worker_id)] = record.timestamp


class HandlerLevelFeedArbHandler(FeedArbitrationHandler):
  EventDistributorType = SwitchingArbEventDistributor

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._handler_map = {}
    self._chosen_ms_handler_wid = None
    self._last_book_ts = {}

  def add_check_product(self, product):
    for handler in self._handler_map.values():
      handler.add_check_product(product)

  def add_handler(self, worker_id, handler_cls):
    self._handler_map[str(worker_id)] = handler_cls(feed_checker_config=self._feed_checker_config,
                                                    event_dist=self._event_dist)

  def post_subscribe(self, *args, **kawrgs):
    for handler in self._handler_map.values():
      handler.post_subscribe(*args, **kawrgs)

  def set_live(self, is_live):
    assert len(self._handler_map) > 0
    for handler in self._handler_map.values():
      handler.set_live(is_live)

  def subscribe_flow(self, flow_sub, topic, callback_attr, hash_function):
    for worker_id, handler in self._handler_map.items():
      callback = getattr(handler, callback_attr)
      wrap_callback = self.get_wrapper_callback(callback, hash_function)
      flow_sub.subscribe('%s*%s' % (topic, worker_id), wrap_callback)

  def subscribe_multi_shift(self, flow_sub, topic, callback_attr):
    for worker_id, handler in self._handler_map.items():
      callback = getattr(handler, callback_attr)
      wrap_callback = self.get_multi_shift_callback(handler, callback)
      flow_sub.subscribe('%s*%s' % (topic, worker_id), wrap_callback)
      if self._chosen_ms_handler_wid is None:
        self._chosen_ms_handler_wid = worker_id

  def get_multi_shift_callback(self, handler, callback):
    def wrap_ms_handler(record, queue_data, topic_data):
      wid = int(get_worker_id(record.topic_id))
      if (record.timestamp - self._last_book_ts.get(int(self._chosen_ms_handler_wid), 0) >
          self._old_feed_timeout):
        self._chosen_ms_handler_wid = wid

      prev_publish_try = self._event_dist.publish_try
      if int(self._chosen_ms_handler_wid) == wid:
        self._event_dist.set_publish(True)
        callback(record, queue_data, topic_data)
        self._event_dist.set_publish(False)
      else:
        callback(record, queue_data, topic_data)
      if self._event_dist.publish_try - prev_publish_try > 0:
        self._last_book_ts[wid] = record.timestamp

    return wrap_ms_handler

  def get_wrapper_callback(self, callback, hash_function):
    def wrap_handler(record, queue_data, topic_data):
      uid = hash_function(record)
      if self.update_and_check_new(record.timestamp, uid):
        self._event_dist.set_publish(True)
        callback(record, queue_data, topic_data)
        self._event_dist.set_publish(False)
      else:
        callback(record, queue_data, topic_data)
      self._last_feed_ts[int(get_worker_id(record.topic_id))] = record.timestamp

    return wrap_handler
